from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.trustedhost import TrustedHostMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn
import time
from complaint_extractor_local import ComplaintProcessorLocal
from content_classifier import ContentClassifier
from typing import List, Dict

app = FastAPI()

# 将 request_counts 移动到全局作用域
request_counts = {}

# 添加安全头部
@app.middleware("http")
async def add_security_headers(request: Request, call_next):
    response = await call_next(request)
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["X-XSS-Protection"] = "1; mode=block"
    return response

# 限制请求频率
@app.middleware("http")
async def rate_limit(request: Request, call_next):
    global request_counts  # 声明使用全局变量
    client_ip = request.client.host
    current_time = time.time()

    # 清理旧的请求记录
    request_counts = {
        ip: count 
        for ip, count in request_counts.items()
        if current_time - count["timestamp"] < 60
    }

    if client_ip in request_counts:
        if request_counts[client_ip]["count"] > 60:  # 每分钟最多60个请求
            return JSONResponse(
                status_code=429,
                content={"error": "Too many requests"}
            )
        request_counts[client_ip]["count"] += 1
    else:
        request_counts[client_ip] = {"count": 1, "timestamp": current_time}

    return await call_next(request)

# 添加 GZIP 压缩
app.add_middleware(GZipMiddleware)

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

class ComplaintRequest(BaseModel):
    complaint_text: str

@app.post("/api/analyze_complaint")
async def analyze_complaint(request: ComplaintRequest):
    processor = ComplaintProcessorLocal()
    result = processor.process_complaint(request.complaint_text)
    print(result)
    if result:
        return {"code": "1", "data": result}
    return {"code": "0", "message": "处理失败"}

# 添加新的请求模型
class CategoryRequest(BaseModel):
    content: str
    categories: List[Dict[str, str]]

# 添加新的接口
@app.post("/api/content/category")
async def classify_content(request: CategoryRequest):
    try:
        classifier = ContentClassifier()
        result = classifier.classify_content(request.content, request.categories)
        
        if result and "分类结果" in result:
            return {
                "code": "1",
                "data": result["分类结果"]
            }
        return {
            "code": "0",
            "message": "分类失败"
        }
    except Exception as e:
        return {
            "code": "0",
            "message": f"处理出错: {str(e)}"
        }

# 静态文件挂载放在最后
app.mount("/", StaticFiles(directory="static", html=True), name="static")

if __name__ == "__main__":
    uvicorn.run(
        "main:app",  # 修改这里，使用字符串形式的导入路径
        host="0.0.0.0",
        port=8000,
        workers=4,
        reload=True  # 添加热重载支持
    )